import glob
import os
import random
import shutil

import cv2
import numpy as np
import yaml

from user_util import cv_imread, read_xml_annotation,cv_imwrite
from user_det_random_clip_sg import clip_muti_dirs

class sengo2yolo:
    '''
    sengo转为v5格式数据集
    '''
    def __init__(self,
                 glob_str=r'D:\data\231215安全带\trainV8Det_flball_blue\_add_imgs',
                 yolov5_dir=r'D:\data\231215安全带\trainV8Det_flball_blue\format_data'):
        self.glob_str = glob_str # source
        self.yolov5_dir = yolov5_dir # dst
        self.target_size = 640

        self.sizerange = [0.1, 0.8]
        self.pos_cnt = 4 # repeat
    def run(self,):
        '''
        sengo随机裁剪，转为v5格式数据集
        Returns:

        '''
        # from common.utils import read_xml_annotation
        # 随机裁剪
        clip_muti_dirs(self.glob_str, self.sizerange, self.pos_cnt)

        # 0 统计芯歌格式标签的类别和个数, 生成data.yaml
        name_count_dict = self.count_xglabel(glob_str=self.glob_str)
        name_id_dict = {}
        names_list = []
        for ind, (k, v) in enumerate(name_count_dict.items()):
            name_id_dict[k] = ind
            names_list.append(k)
        print(name_id_dict)
        with open(self.yolov5_dir + '\classes.txt', mode="w", encoding="utf-8") as f:
            f.write(str(name_id_dict))
        #
        data_yaml_path = f'{self.yolov5_dir}/data.yaml'
        dataset_dict = {
            'train': f'{self.yolov5_dir}/images/train',
            'val': f'{self.yolov5_dir}/images/val',
            'nc': len(names_list),
            'names': names_list
        }
        with open(data_yaml_path, 'w', encoding='utf-8') as f:
            yaml.dump(data=dataset_dict, stream=f, allow_unicode=True)  # 保存

        # 1 get ls_xml, ls_jpg
        print('1 get ls_xml, ls_jpg')
        ls_xml = []
        ls_jpg = []
        for root, dirs, files in os.walk(self.glob_str):  # os.walk 递归遍历
            for fn in files:
                fn0, ext = os.path.splitext(fn)
                if ext.lower() in ['.xml'] :
                    fn1 = os.path.join(root, fn)
                    ls_xml.append(fn1)
                # if ext.lower() in ['.jpg', '.jpeg'] and '_mini' not in root:
                if ext.lower() in ['.jpg', '.jpeg']:
                    fn1 = os.path.join(root, fn)
                    ls_jpg.append(fn1)
        print(f'ls_xml {len(ls_xml)}')
        print(f'ls_jpg {len(ls_jpg)}')

        # 2 xml 转为 txt
        print('2 xml 转为 txt')
        for i in ls_xml:
            xywhnames = read_xml_annotation(i, divwh=True)  # [[xywh], name ]
            txt_str = ''
            for [x, y, w, h], name in xywhnames:
                txt_str += f'{name_id_dict[name]} {x} {y} {w} {h}\n'
            # save
            save_txt_name = i[:-3] + 'txt'
            # print(save_txt_name)
            with open(save_txt_name, mode="w", encoding="utf-8") as f:
                f.write(txt_str[:-1])

        # 3 get dataset_yolov5
        print('3 get dataset_yolov5')
        self.update_yolov5(self.yolov5_dir, ls_jpg)

        # 4 resize
        print('4 resize')
        self.resize_clip_imgs(self.yolov5_dir, target_size=self.target_size)



    def count_xglabel(self, glob_str=''):
        '''
        统计芯歌格式标签的类别和个数
        20240201
        Returns:

        '''

        # ls = glob.glob(glob_str)
        exts = ['.xml', ]
        ls = []  # xml文件地址list
        ls_jpg = []
        for root, dirs, files in os.walk(glob_str):  # os.walk 递归遍历
            for fn in files:
                fn0, ext = os.path.splitext(fn)
                if ext.lower() in exts and '_mini' not in root:  # 非mini
                    fn1 = os.path.join(root, fn)
                    ls.append(fn1)
                if ext.lower() in ['.jpg', '.jpeg'] and '_mini' not in root:  # 非mini
                    fn1 = os.path.join(root, fn)
                    ls_jpg.append(fn1)
        print(f'ls_xml {len(ls)}')
        print(f'ls_jpg {len(ls_jpg)}')

        res = {}  # clsname: cnt
        for i in ls:
            xyxy_clsnames = read_xml_annotation(i)
            for xyxy, clsname in xyxy_clsnames:
                if not res.get(clsname):
                    res[clsname] = 1
                else:
                    res[clsname] += 1

        import pprint
        print(f'Folder: \n{glob_str}')
        print('Number of each defect type: ')
        pprint.pprint(res)

        cnt = 0
        for k, v in res.items():
            cnt += v
        print(cnt)
        return res

    def update_yolov5(self, yolov5_dir, ls_jpg):
        '''
        make yolov5 or add a unit to yolov5
        '''

        # dst
        if not os.path.exists(yolov5_dir):
            os.makedirs(yolov5_dir)
        else:
            print(f'rm {yolov5_dir}/*/*...')
            if os.path.exists(f'{yolov5_dir}/images'):
                shutil.rmtree(f'{yolov5_dir}/images')  # 删除，重新转换,目前不能增量式添加数据
            if os.path.exists(f'{yolov5_dir}/labels'):
                shutil.rmtree(f'{yolov5_dir}/labels')
        out_fir_tra_img = yolov5_dir + '/images/train/'
        out_fir_val_img = yolov5_dir + '/images/val/'
        out_fir_tra_txt = yolov5_dir + '/labels/train/'
        out_fir_val_txt = yolov5_dir + '/labels/val/'
        iltv_dir = [out_fir_tra_img, out_fir_val_img, out_fir_tra_txt, out_fir_val_txt]
        for i in iltv_dir:
            if not os.path.exists(i):  # 如果不存在
                os.makedirs(i)

        print('before:')
        t_i, v_i, t_tx, v_tx = len(os.listdir(out_fir_tra_img)), len(os.listdir(out_fir_val_img)), \
            len(os.listdir(out_fir_tra_txt)), len(os.listdir(out_fir_val_txt))
        print(f'{out_fir_tra_img}:{t_i}')
        print(f'{out_fir_val_img}:{v_i}')
        print(f'{out_fir_tra_txt}:{t_tx}')
        print(f'{out_fir_val_txt}:{v_tx}')

        # 划分 各类随机抽1%
        # ls_img = glob.glob(self.jpg)
        ls_img = ls_jpg
        random.shuffle(ls_img)
        flag = max(int(0.2 * len(ls_img)) + 1, 1)  # 20%
        # ls_img_val = ls_img[:flag]
        # ls_img_tra = ls_img[flag:]
        ls_img_val = ls_img[:flag]
        ls_img_tra = ls_img[0:]

        # 添加至验证集

        for ind, i in enumerate(ls_img_val):
            # print(f'{ind}/ {ls_img_val}')
            # jpg
            # nam = i.split('/')[-1]
            # nam = i.split('\\')[-1]  # windows
            nam = os.path.basename(i)
            # self.buff = i.split('\\')[-2]  # user
            self.buff = os.path.basename(os.path.dirname(i))
            shutil.copy(i, out_fir_val_img + f'{self.buff}_{nam}')  # 添加至val_img
            # lab
            # txt_abpath = self.txt + nam[:-4] + '.txt'  #
            txt_abpath = f'{i[:-4]}.txt'  #
            if os.path.exists(txt_abpath):  # 如果有标签文件
                shutil.copy(txt_abpath, out_fir_val_txt + f'{self.buff}_{nam[:-4]}.txt')

        # 添加至训练集
        for ind, i in enumerate(ls_img_tra):
            # jpg
            # nam = i.split('/')[-1]
            # nam = i.split('\\')[-1]  # windows
            nam = os.path.basename(i)
            # self.buff = i.split('\\')[-2]  # user 父文件名为前缀
            self.buff = os.path.basename(os.path.dirname(i))
            shutil.copy(i, out_fir_tra_img + f'{self.buff}_{nam}')  # 添加至tra_img
            # lab
            txt_abpath = f'{i[:-4]}.txt'
            if os.path.exists(txt_abpath):
                shutil.copy(txt_abpath, out_fir_tra_txt + f'{self.buff}_{nam[:-4]}.txt')  # 添加至tra_lab

        # check
        print('after:')
        t_i_d, v_i_d, t_tx_d, v_tx_d = len(os.listdir(out_fir_tra_img)), len(os.listdir(out_fir_val_img)), \
            len(os.listdir(out_fir_tra_txt)), len(os.listdir(out_fir_val_txt))
        print(f'{out_fir_tra_img}:{t_i_d}')
        print(f'{out_fir_val_img}:{v_i_d}')
        print(f'img added: {t_i_d + v_i_d - t_i - v_i}')
        print(f'{out_fir_tra_txt}:{t_tx_d}')
        print(f'{out_fir_val_txt}:{v_tx_d}')
        print(f'lab added: {t_tx_d + v_tx_d - t_tx - v_tx}')

        print(f'Update_yolov5, Done.')

    def resize_clip_imgs(self, yolov5_dir, target_size=416):  # 缩小训练集图片分辨率，加快训练

        imgs_glob = yolov5_dir + '/images/*/*.jpg'
        ls = glob.glob(imgs_glob)
        scale = 1.0
        for i in ls:
            img = cv_imread(i)
            scale = target_size / max(img.shape[:2])
            # print(scale)
            try:
                # img_split1 = cv2.resize(img_split1, dsize=None, fx=scale, fy=scale)
                # img_split2 = cv2.resize(img_split2, dsize=None, fx=scale, fy=scale)
                img = cv2.resize(img, dsize=None, fx=scale, fy=scale)
            except:
                print(i)
            # cv2.imwrite(i[:-4]+'_split1.jpg', img_split1)
            # cv2.imwrite(i[:-4] + '_split2.jpg', img_split2)
            cv_imwrite(i, img)
        print(f'resized done, total {len(ls)}, scale {scale}')

if __name__ == '__main__':
    s2y = sengo2yolo(glob_str=r'D:\data\231215安全带\trainV8Det_flball_blue\_add_imgs',
                     yolov5_dir=r'D:\data\231215安全带\trainV8Det_flball_blue\format_data')
    s2y.run()